import logging
import sys
import time
import websocket

from threading import Thread

from stomppy import frame

VERSIONS = '1.0,1.1'
CONNECT_RETRY_SLEEP = 1
HEART_BEAT = '10000,10000'

CMD_ACK = 'ACK'
CMD_CONNECT = 'CONNECT'
CMD_CONNECTED = 'CONNECTED'
CMD_DISCONNECT = 'DISCONNECT'
CMD_ERROR = 'ERROR'
CMD_MESSAGE = 'MESSAGE'
CMD_NACK = 'NACK'
CMD_RECEIPT = 'RECEIPT'
CMD_SEND = 'SEND'
CMD_SUBSCRIBE = 'SUBSCRIBE'
CMD_UNSUBSCRIBE = 'UNSUBSCRIBE'

class StompClient:
    def __init__(self, url):
        self.url = url
        self.ws = websocket.WebSocketApp(self.url)
        self.ws.on_open = self._on_open
        self.ws.on_message = self._on_message
        self.ws.on_error = self._on_error
        self.ws.on_close = self._on_close

        self.opened = False
        self.connected = False
        self.counter = 0
        self.subscriptions = {}
        self._connectCallback = None
        self._errorCallback = None

    def connect(self, 
                login='',
                passcode='',
                connectCallback=None,
                errorCallback=None,
                headers={},
                timeout=0):
        logging.debug(f'Try connect to {self.url}')

        self._connect(timeout)

        headers['host'] = self.url
        headers['accept-version'] = VERSIONS
        headers['heart-beat'] = HEART_BEAT

        if login:
            headers['login'] = login
        if passcode:
            headers['passcode'] = passcode

        self._connectCallback = connectCallback
        self._errorCallback = errorCallback

        self._transmit(CMD_CONNECT, headers)

    def disconnect(self, disconnectCallback=None, headers={}):
        logging.debug(f'Disconnect from {self.url}')

        self._transmit(CMD_DISCONNECT, headers)
        self.ws.close()
        self._clean_up()

        if disconnectCallback:
            disconnectCallback()

    def send(self, destination, body='', headers={}):
        headers['destination'] = destination
        self._transmit(CMD_SEND, headers, body)

    def subscribe(self, destination, callback=None, headers={}):
        if 'id' not in headers:
            headers['id'] = 'sub-' + str(self.counter)
            self.counter += 1
        headers['destination'] = destination
        self.subscriptions[headers['id']] = callback
        self._transmit(CMD_SUBSCRIBE, headers)
        
        return headers['id']

    def unsubscribe(self, sub_id):
        del self.subscriptions[sub_id]
        return self._transmit(CMD_UNSUBSCRIBE, {'id': sub_id})

    def ack(self, message_id, subscription, headers={}):
        headers['meesage-id'] = message_id
        headers['subscription']= subscription
        return self._transmit(CMD_ACK, headers)

    def nack(self, message_id, subscription, headers={}):
        headers['meesage-id'] = message_id
        headers['subscription']= subscription
        return self._transmit(CMD_NACK, headers)

    def _connect(self, timeout=0): # timeout is ms
        thread = Thread(target=self.ws.run_forever)
        thread.daemon = True
        thread.start()

        time_used_ms = 0
        while not self.opened:
            time.sleep(CONNECT_RETRY_SLEEP)
            time_used_ms = CONNECT_RETRY_SLEEP * 1000
            if (timeout > 0) and (time_used_ms > timeout):
                raise TimeoutError(f'Connection to {self.url} timed out')

    def _on_open(self, ws):
        logging.debug('on_open successfully')
        self.opened = True

    def _on_message(self, ws, message):
        logging.debug('on_message: \n' + str(message))

        f = frame.Frame.unmarshall(message)
        _results = []
        if f.command == CMD_CONNECTED:
            self.connected = True
            logging.debug(f'Connected to server {self.url}')
            if self._connectCallback:
                r = self._connectCallback(f)
                _results.append(r)
        elif f.command == CMD_MESSAGE:
            subscription = f.headers['subscription']
            if subscription in self.subscriptions:
                on_receive = self.subscriptions[subscription]
                message_id = f.headers['message-id']

                def ack(headers):
                    return self.ack(message_id, subscription, headers)

                def nack(headers):
                    return self.nack(message_id, subscription, headers)

                f.ack = ack
                f.nack = nack

                _results.append(on_receive(f))
            else:
                info = 'Unhandled received message: ' + str(f)
                logging.debug(info)
                _results.append(info)
        elif f.command == CMD_RECEIPT:
            pass
        elif f.command == CMD_ERROR:
            if self._errorCallback:
                _results.append(self._errorCallback(f))
        else:
            info = 'Unknown command: ' + f.command
            logging.debug(info)
            _results.append(info)

        return _results

    def _on_error(self, ws, error):
        logging.debug('on_error: ' + str(error))

    def _on_close(self, ws, status, message):
        self._clean_up()
        logging.debug(f'Connection to {self.url} lost')
        
    def _clean_up(self):
        self.connected = False

    def _transmit(self, command, headers={}, body=''):
        data = frame.Frame.marshall(command, headers, body)
        logging.debug('_transmit:\n' + data)
        self.ws.send(data)


# Study from https://github.com/GlassyWing/stomp_ws_py